package com.example.jucsample.threadpool;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author zhenwei.wu
 * @date 2021/2/13
 */
public class MyThreadPool implements Executor {

    private int corePoolSize;
    private int maximumPoolSize;
    private BlockingQueue<Runnable> queue;
    //记录当前工作线程数
    private AtomicInteger count;
    private long keepAliveTime;
    private RejectHandler rejectHandler;

    public MyThreadPool(int corePoolSize, int maximumPoolSize, BlockingQueue<Runnable> queue, long keepAliveTime, RejectHandler rejectHandler) {
        this.corePoolSize = corePoolSize;
        this.maximumPoolSize = maximumPoolSize;
        this.queue = queue;
        this.keepAliveTime = keepAliveTime;
        this.rejectHandler = rejectHandler;
        count = new AtomicInteger(0);
    }

    @Override
    public void execute(Runnable task) {
        int ct = count.get();
        //核心线程数未满，尝试增加核心线程
        if (ct < corePoolSize && count.compareAndSet(ct, ct + 1)) {
            new Worker(task).start();
            return;
        }
        //入队
        if (queue.offer(task)) {
            return;
        }
        //重新获取一遍count，否则如果在core分支cas失败，此处必然也失败
        ct = count.get();
        //队列已满，尝试增加非核心线程
        if (ct < maximumPoolSize && count.compareAndSet(ct, ct + 1)) {
            new Worker(task).start();
            return;
        }
        //已达最大线程数，拒绝
        rejectHandler.reject(task);

    }

    class Worker extends Thread {

        Runnable firstTask;

        public Worker(Runnable firstTask) {
            this.firstTask = firstTask;
        }

        @Override
        public void run() {
            Runnable task = firstTask;
            firstTask = null;
            while (true) {
                try {
                    //getTask会阻塞
                    if (task != null || (task = getTask()) != null) {
                        task.run();
                    } else {
                        //getTask超时才会进入，直接退出，线程销毁
                        break;
                    }
                } catch (InterruptedException e) {
                    e.printStackTrace();
                } finally {
                    //置空，否则不能getTask
                    task = null;
                }
            }
        }

    }

    Runnable getTask() throws InterruptedException {
        //标记是否超时过
        boolean timedOut = false;
        while (true) {
            int ct = count.get();
            //超出核心线程数才进入超时逻辑，即使timeOut由于线程poll超时过一次变成true，执行到这里如果不超出corePoolSize，可以再次进入take分支
            if (ct > corePoolSize) {
                //超出核心线程数
                if (timedOut) {
                    //已超时过，尝试减少工作线程数，失败会continue，然后重新比较corePoolSize，重试减少线程数
                    if (count.compareAndSet(ct, ct - 1)) {
                        return null;
                    } else {
                        continue;
                    }
                }
                Runnable task = queue.poll(keepAliveTime, TimeUnit.MILLISECONDS);
                if (task == null) {
                    //poll超时才进入
                    timedOut = true;
                    continue;
                }
                return task;
            } else {
                //必然能获取到task
                return queue.take();
            }
        }
    }

    public static interface RejectHandler {

        void reject(Runnable r);

    }


    public static void main(String[] args) {
        MyThreadPool pool = new MyThreadPool(2, 5, new LinkedBlockingQueue<>(1), 2000, r -> {
            System.out.println(r + ": reject");
        });
        for (int i = 0; i < 3; i++) {
            final int x = i;
            new Thread(() -> {
                for (int j = 0; j < 5; j++) {
                    final int y = j;
                    pool.execute(() -> {
                        try {
                            Thread.sleep(3000L);
                        } catch (InterruptedException e) {
                            e.printStackTrace();
                        }
                        LocalDateTime now = LocalDateTime.now();
                        System.out.println(String.format("线程i=%s, j=%s,执行结束: %s", x, y, now.format(DateTimeFormatter.ISO_DATE_TIME)));
                    });
                }
            }).start();

        }

    }

}
